import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

class ImgProcessing(object):
    def __init__(self,img,model_path='static/moban.jpg'):
        self.img = img
        self.img_model = cv.imread(model_path)
        self.model_circle,self.model_r = self.circle(self.img_model)
        self.img_binary_image = self.img_proccess(self.img)
        self.model_binary_image = self.img_proccess(self.img_model)
        self.img_end = self.transformation()
        self.img_circle,self.img_r = self.circle(self.img_end)

    # 图片预处理（降噪、灰度、阈值切割）
    def img_proccess(self,img):
        # 降噪
        img_gaosi = cv.GaussianBlur(img, (5, 5), 0)
        # img_median = cv.medianBlur(img_gray,5)
        # img_billateral = cv.bilateralFilter(img_gray,9,75,75)
        # 转换为灰度图
        img_gray = cv.cvtColor(img_gaosi, cv.COLOR_BGR2GRAY)
        # 阈值切割
        ret, binary_image = cv.threshold(img_gray, 200, 255, cv.THRESH_BINARY)
        return binary_image

    # 获取圆心
    def circle(self,img):
        img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        img_gaosi = cv.GaussianBlur(img_gray, (5, 5), 0)
        # 霍夫圆检测
        circles = cv.HoughCircles(
            img_gaosi,
            cv.HOUGH_GRADIENT,
            dp=1,  # 累加器分辨率与图像分辨率的反比
            minDist=100,  # 检测到的圆心之间的最小距离
            param1=100,  # Canny边缘检测的高阈值
            param2=50,  # 累加器阈值，越小检测到的圆越多
            minRadius=10,  # 最小半径
            maxRadius=100  # 最大半径，0表示不限制
        )
        circles_num = []
        r = []
        if circles is not None:
            circles = np.uint16(np.around(circles))

            for i, circle in enumerate(circles[0,:]):
                circle_num = [int(circle[0]), int(circle[1])]
                circles_num.append(circle_num)
                r.append(circle[2])
        else:
            print("未检测到圆")

        return circles_num,r

    # 计算欧式距离
    def distance(self):
        self.model_circle = np.array(self.model_circle)
        self.img_circle = np.array(self.img_circle)
        sorted_indices = np.lexsort((self.model_circle[:, 0], self.model_circle[:, 1]))  # 先按arr[:,1]，再按arr[:,0]
        sorted_arr = self.model_circle[sorted_indices]

        sorted_indices = np.lexsort((self.img_circle[:, 0], self.img_circle[:, 1]))  # 先按arr[:,1]，再按arr[:,0]
        img_sorted_arr = self.img_circle[sorted_indices]

        dis1 = np.linalg.norm(sorted_arr[0] - img_sorted_arr[0])
        dis2 = np.linalg.norm(sorted_arr[1] - img_sorted_arr[1])
        dis3 = np.linalg.norm(sorted_arr[2] - img_sorted_arr[2])
        dis4 = np.linalg.norm(sorted_arr[3] - img_sorted_arr[3])
        return np.array([dis1, dis2, dis3, dis4])

    # 透视变换
    def transformation(self):
        shift = cv.xfeatures2d.SIFT_create()
        kd1, dst1 = shift.detectAndCompute(self.model_binary_image, None)
        kd2, dst2 = shift.detectAndCompute(self.img_binary_image, None)
        FLANN_INDEX_KDTREE = 1
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        flann = cv.FlannBasedMatcher(index_params, search_params)

        matches = flann.knnMatch(dst1, dst2, k=2)

        goods = []
        for m, n in matches:
            if m.distance < 0.9 * n.distance:
                goods.append(m)

        src_pts = np.float32([kd1[m.queryIdx].pt for m in goods]).reshape(-1, 1, 2)
        dst_pts = np.float32([kd2[m.trainIdx].pt for m in goods]).reshape(-1, 1, 2)
        M, mask = cv.findHomography(dst_pts, src_pts, cv.RANSAC, 5.0)
        width, height = self.img_model.shape[:2]
        img_end = cv.warpPerspective(self.img, M, (width, height))
        return img_end

    def img_polt(self):
        i = 0
        for x,y in self.img_circle:
            cv.circle(self.img_end, (x, y), self.img_r[i], (0, 255, 0), 2)
            # 绘制圆心
            cv.circle(self.img_end, (x, y), 2, (0, 0, 255), 3)
            i += 1
    def img_show(self,path):
        self.img_polt()
        cv.imwrite(path,self.img_end)